package mr

import (
	//"go/ast"
	"log"
	"sync"
	"time"
)
import "net"
import "os"
import "net/rpc"
import "net/http"

type TaskStatus int

const (
	idle TaskStatus = iota
	running
	finished
	failed
)

type MapTaskRequest struct {
	TaskId    int
	Status    TaskStatus
	StartTime int64
}

type ReduceTaskRequest struct {
	Status    TaskStatus
	StartTime int64
}

type Coordinator struct {
	// Your definitions here.
	NReduce       int
	MapTasks      map[string]*MapTaskRequest
	mapSuccess    bool
	mutex_map     sync.Mutex // 细粒度的锁，对map函数上锁
	ReduceTasks   []*ReduceTaskRequest
	reduceSuccess bool
	mutex_reduce  sync.Mutex // 对reduce函数上锁
}

// Your code here -- RPC handlers for the worker to call.
// 初始化所有的 map 和 reduce 任务
func (c *Coordinator) initTask(files []string) {
	for idx, fileName := range files {
		c.MapTasks[fileName] = &MapTaskRequest{
			TaskId: idx,
			Status: idle,
		}
	}

	for idx := range c.ReduceTasks {
		c.ReduceTasks[idx] = &ReduceTaskRequest{
			Status: idle,
		}
	}
}

// RPC的响应函数
func (c *Coordinator) AskForTask(req *MessageSend, reply MessageReply) error {
	if req.MsgType != AskForTask {
		return BadMsgType
	}
	if !c.mapSuccess {

		c.mutex_map.Lock()

		count_map_success := 0
		for filename, task := range c.MapTasks {
			assign := false

			if task.Status == idle || task.Status == failed {
				assign = true
			} else if task.Status == running {
				cur_time := time.Now().Unix()
				if cur_time-task.StartTime > 10 {
					task.StartTime = cur_time
					assign = true
				}
			} else {
				count_map_success++
			}

			// 将未分配的任务和已经失败的任务分配给这个worker
			if assign {
				reply.MsgType = MapTaskAlloc
				reply.TaskName = filename
				reply.NReduce = c.NReduce
				reply.TaskId = task.TaskId

				task.Status = running
				task.StartTime = time.Now().Unix()
				c.mutex_map.Unlock()
				return nil
			}
		}

		// map任务没有可以分配的, 但都还未完成
		if count_map_success < len(c.MapTasks) {
			reply.MsgType = Wait
			return nil
		} else {
			c.mapSuccess = true
		}
	}

	// 运行到这里说明map任务都已经完成
	if !c.reduceSuccess {
		c.mutex_reduce.Lock()

		count_reduce_success := 0
		for idx, task := range c.ReduceTasks {
			assign := false
			if task.Status == idle || task.Status == failed {
				assign = true
			} else if task.Status == running {
				cur_time := time.Now().Unix()
				if cur_time-task.StartTime > 10 {
					task.StartTime = cur_time
					assign = true
				}
			} else {
				count_reduce_success++
			}

			if assign {
				reply.MsgType = ReduceTaskAlloc
				reply.TaskId = idx

				task.Status = running
				task.StartTime = time.Now().Unix()

				c.mutex_reduce.Unlock()
			}
		}

		c.mutex_reduce.Unlock()

		if count_reduce_success < len(c.ReduceTasks) {
			reply.MsgType = Wait
			return nil
		} else {
			c.reduceSuccess = true
		}
	}

	reply.MsgType = Shutdown

	return nil
}

// RPC响应函数
func (c *Coordinator) NoticeResult(req *MessageSend, reply *MessageSend) error {
	if req.MsgType == MapSuccess {
		c.mutex_map.Lock()

		for _, t := range c.MapTasks {
			if t.TaskId == req.TaskId {
				t.Status = finished

				c.mutex_map.Unlock()

				return nil
			}
		}

		c.mutex_map.Unlock()
	} else if req.MsgType == ReduceSuccess {
		c.mutex_reduce.Lock()

		c.ReduceTasks[req.TaskId].Status = finished

		c.mutex_reduce.Unlock()
		return nil
	} else if req.MsgType == MapFailed {
		c.mutex_map.Lock()
		for _, t := range c.MapTasks {
			if t.TaskId == req.TaskId && t.Status == running {
				t.Status = failed

				c.mutex_map.Unlock()
				return nil
			}
		}
		c.mutex_map.Unlock()
	} else if req.MsgType == ReduceFailed {
		c.mutex_reduce.Lock()

		if c.ReduceTasks[req.TaskId].Status == running {
			c.ReduceTasks[req.TaskId].Status = failed
		}
		c.mutex_reduce.Unlock()
		return nil
	}

	return nil
}

// an example RPC handler.
//
// the RPC argument and reply types are defined in rpc.go.
func (c *Coordinator) Example(args *ExampleArgs, reply *ExampleReply) error {
	reply.Y = args.X + 1
	return nil
}

// start a thread that listens for RPCs from worker.go
func (c *Coordinator) server() {
	rpc.Register(c)
	rpc.HandleHTTP()
	//l, e := net.Listen("tcp", ":1234")
	sockname := coordinatorSock()
	os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
}

// main/mrcoordinator.go calls Done() periodically to find out
// if the entire job has finished.
func (c *Coordinator) Done() bool {

	// Your code here.
	for _, task := range c.MapTasks {
		if task.Status != finished {
			return false
		}
	}

	for _, task := range c.ReduceTasks {
		if task.Status != finished {
			return false
		}
	}

	return true
}

// create a Coordinator.
// main/mrcoordinator.go calls this function.
// nReduce is the number of reduce tasks to use.
func MakeCoordinator(files []string, nReduce int) *Coordinator {

	// Your code here.
	c := Coordinator{
		NReduce:     nReduce,
		MapTasks:    make(map[string]*MapTaskRequest),
		ReduceTasks: make([]*ReduceTaskRequest, nReduce),
	}

	c.initTask(files)

	c.server()
	return &c
}
